source activate base # autodl
# source activate ke2torch23cu121 # 3090
export HUGGINGFACE_CACHE=/root/autodl-fs/huggingface/
export PYTHONUNBUFFERED=1

model=llama2-7b-chat
method=ICE

DATE=$(date +"%Y-%m-%d")
# autodl
if [ ! -d "/root/autodl-fs/${method}-log/${DATE}" ]; then
    mkdir -p /root/autodl-fs/${method}-log/${DATE}
    mkdir -p /root/autodl-fs/${method}-log/${DATE}/${method}-js
    mkdir -p /root/autodl-fs/${method}-log/${DATE}/${method}-kl
fi
# # 3090
# if [ ! -d "./${method}-log/${DATE}" ]; then
#     mkdir -p ./${method}-log/${DATE}
#     mkdir -p ./${method}-log/${DATE}/${method}-js
#     mkdir -p ./${method}-log/${DATE}/${method}-kl
# fi

# # KnowEdit datasets
# declare -A datasets=(
#     ["zsre"]="../dataset-ke/KnowEdit/benchmark/ZsRE/ZsRE-test-all.json"
#     ["recent"]="../dataset-ke/KnowEdit/benchmark/wiki_recent/recent_test.json"
#     ["counterfact"]="../dataset-ke/KnowEdit/benchmark/wiki_counterfact/test_cf.json"
#     ["wikibio"]="../dataset-ke/KnowEdit/benchmark/WikiBio/wikibio-test-all.json"
# )

gpu_id=0
cnt=gen-pre-edit
# ICE KnowEdit datasets
declare -A datasets=(
    ["zsre"]="../dataset-ke/${method}/zsre.json"
    # ["recent"]="../dataset-ke/${method}/wikidata_recent.json"
    ["counterfact"]="../dataset-ke/${method}/wikidata_counterfact.json"
    # ["wikibio"]="../dataset-ke/${method}/wikibio.json"
)
# --test_generation 是否测试生成,计算fluency指标
# --save_gen_sentence 是否保存生成的句子,观察编辑后输出重复等问题
# 默认为false,指定位true
# for method_type in js kl; do
for method_type in kl; do
    for datatype in "${!datasets[@]}"; do
        data_dir="${datasets[$datatype]}"
        echo "${DATE}/${method}-${method_type}/${model}-${datatype}-${gpu_id}"
        CUDA_VISIBLE_DEVICES=${gpu_id} python examples/run_knowedit_llama2.py \
            --editing_method=${method} \
            --hparams_dir=./hparams/${method}/${model}-${method_type}.yaml \
            --data_dir=${data_dir} \
            --datatype=${datatype} \
            --test_generation \
            --save_gen_sentence \
            --metrics_save_dir=/root/autodl-fs/${method}-results/${model}/${method}-${method_type}-${DATE}-${cnt} \
            --gen_sentence_save_dir=/root/autodl-fs/${method}-outputs/${model}/${method}-${method_type}-${DATE}-${cnt} \
            --pre_file=/root/autodl-fs/pre_edit_autodl/${model}_${datatype}_pre_edit.json \
            > /root/autodl-fs/${method}-log/${DATE}/${method}-${method_type}/${model}-${datatype}-${gpu_id}-${cnt}.log 2>&1 &
        wait $!
    done &
done &

gpu_id=1
declare -A datasets=(
    # ["zsre"]="../dataset-ke/${method}/zsre.json"
    ["recent"]="../dataset-ke/${method}/wikidata_recent.json"
    # ["counterfact"]="../dataset-ke/${method}/wikidata_counterfact.json"
    ["wikibio"]="../dataset-ke/${method}/wikibio.json"
)
for method_type in kl; do
    for datatype in "${!datasets[@]}"; do
        data_dir="${datasets[$datatype]}"
        echo "${DATE}/${method}-${method_type}/${model}-${datatype}-${gpu_id}"
        CUDA_VISIBLE_DEVICES=${gpu_id} python examples/run_knowedit_llama2.py \
            --editing_method=${method} \
            --hparams_dir=./hparams/${method}/${model}-${method_type}.yaml \
            --data_dir=${data_dir} \
            --datatype=${datatype} \
            --test_generation \
            --save_gen_sentence \
            --metrics_save_dir=/root/autodl-fs/${method}-results/${model}/${method}-${method_type}-${DATE}-${cnt} \
            --gen_sentence_save_dir=/root/autodl-fs/${method}-outputs/${model}/${method}-${method_type}-${DATE}-${cnt} \
            --pre_file=/root/autodl-fs/pre_edit_autodl/${model}_${datatype}_pre_edit.json \
            > /root/autodl-fs/${method}-log/${DATE}/${method}-${method_type}/${model}-${datatype}-${gpu_id}-${cnt}.log 2>&1 &
        wait $!
    done &
done &

wait
# wait && /usr/bin/shutdown

